# model for image classification for jd item pics
from mxnet.gluon.model_zoo import vision
import cv2
import mxnet as mx
import mxnet.ndarray as nd
from data_loader import read_img
from utils import show_img
from mxnet import image
import mxnet.gluon.nn as nn


def load_img(img_file):
    img, _, _ = read_img(img_file)

    img = img / 255.0
    print(img.shape)
    img = img.reshape((224, 224, 3))
    show_img(img_file)

    normalized = mx.image.color_normalize(nd.array(img),
                                          mean=mx.nd.array([0.485, 0.456, 0.406]),
                                          std=mx.nd.array([0.229, 0.224, 0.225]))
    print(normalized)
    return normalized.reshape((3, 224, 224)).expand_dims(0)


def transform(data):
    data = data.transpose((2, 0, 1)).expand_dims(axis=0)
    rgb_mean = nd.array([0.485, 0.456, 0.406]).reshape((1, 3, 1, 1))
    rgb_std = nd.array([0.229, 0.224, 0.225]).reshape((1, 3, 1, 1))
    return (data.astype('float32') / 255 - rgb_mean) / rgb_std


def read_label_file(label_file):
    labels = []
    with open(label_file) as fd:
        for i, l in enumerate(fd):
            d = l.split(' ', 1)[-1]
            labels.append(d)
    return labels


class JDClassifier(nn.Block):
    def __init__(self):
        super(JDClassifier, self).__init__()
        with self.name_scope():
            self.features = nn.HybridSequential()

        with self.name_scope():
            self.output = nn.HybridSequential()

            # self.output.add(nn.Conv2D())

    def forward(self):
        pass


if __name__ == '__main__':
    import sys

    label = read_label_file('./synset.txt')

    image_file = sys.argv[1]

    # input shape Bx3xHxW
    model = vision.densenet121(pretrained=True)
    # data = load_img('/home/alpha/Pictures/640.jpeg')
    x = image.imread(image_file)
    x, _ = image.center_crop(x, (224, 224))
    x = transform(x)
    out = model(x)

    print(nd.max(out))
    preds = [int(i) for i in nd.argsort(out, axis=1, is_ascend=0)[0][0:10].asnumpy().tolist()]
    for c in preds:
        print(label[c])
